import sys
import pandas as pd
import talib
import efinance as ef
import numpy as np
import json


def calculate_kdj(df, n=9, m=3):
    df = df.copy()
    
    # 确保索引是连续整数
    df.reset_index(drop=True, inplace=True)
    
    # 计算 n 日最低价和最高价
    df['low_min'] = df['最低'].rolling(n, min_periods=1).min()
    df['high_max'] = df['最高'].rolling(n, min_periods=1).max()
    
    # 计算 RSV（避免除零错误）
    df['RSV'] = 100 * (df['收盘'] - df['low_min']) / (df['high_max'] - df['low_min'] + 1e-9)
    
    # 初始化 K 和 D 列
    df['K'] = 50.0
    df['D'] = 50.0
    
    # 使用 iloc 进行位置索引计算
    for i in range(1, len(df)):
        rsv_i = df['RSV'].iloc[i]
        
        # 如果 RSV 是有效值
        if pd.notna(rsv_i):
            # 获取前一日的 K 和 D 值
            prev_k = df['K'].iloc[i-1]
            prev_d = df['D'].iloc[i-1]
            
            # 计算当前 K 值
            current_k = (prev_k * (m-1) + rsv_i) / m
            
            # 计算当前 D 值
            current_d = (prev_d * (m-1) + current_k) / m
            
            # 使用 iloc 设置当前值
            df['K'].iloc[i] = current_k
            df['D'].iloc[i] = current_d
    
    # 计算 J 值
    df['J'] = 3 * df['K'] - 2 * df['D']
    
    # 清理临时列
    df.drop(['low_min', 'high_max', 'RSV'], axis=1, inplace=True)
    
    return df


def calculate_indicators(df):
    df = df.sort_values('日期', ascending=True).copy()
    df.reset_index(drop=True, inplace=True)
# ======================================
# 1. 计算 MACD 指标
# ======================================
    fast_period=12
    slow_period=26
    signal_period=9

    df = df.copy()
    
    # 计算快速EMA
    df['EMA_fast'] = df['收盘'].ewm(
        span=fast_period, 
        adjust=False,  # 同花顺使用adjust=False的EMA计算
        min_periods=0  # 允许最小周期为0
    ).mean()
    
    # 计算慢速EMA
    df['EMA_slow'] = df['收盘'].ewm(
        span=slow_period, 
        adjust=False,
        min_periods=0
    ).mean()
    
    # 计算DIF
    df['DIF'] = df['EMA_fast'] - df['EMA_slow']
    
    # 计算DEA (信号线)
    df['DEA'] = df['DIF'].ewm(
        span=signal_period, 
        adjust=False,
        min_periods=0
    ).mean()
    
    # 计算MACD柱状图 (Histogram)
    df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    
    # 清理临时列
    df.drop(['EMA_fast', 'EMA_slow'], axis=1, inplace=True)
# ======================================
# 2. 计算 RSI 指标 (默认周期为14)
# ======================================
    rsi_14 = talib.RSI(df['收盘'], timeperiod=14)
    df['RSI'] = rsi_14

    # ======================================
# 3. 计算 KDJ 指标（随机指标 STOCH 的变体）
# ======================================
    df = calculate_kdj(df, n=9, m=3)
# 4. 计算 布林带 指标（随机指标 STOCH 的变体）
    df['MA20'] = df['收盘'].rolling(window=20).mean()
    df['MA5'] = df['收盘'].rolling(window=5).mean()
    df['MA10'] = df['收盘'].rolling(window=10).mean()
    df['MA30'] = df['收盘'].rolling(window=30).mean()
    df['STD20'] = df['收盘'].rolling(window=20).std()
    df['上轨'] = df['MA20'] + 2 * df['STD20']
    df['下轨'] = df['MA20'] - 2 * df['STD20']
    return df.reset_index()

def get_stock_info_json(stock_code):
    """获取股票信息并返回Python字典"""
    info = ef.stock.get_base_info(stock_code)
    
    # 转换为字典
    if isinstance(info, pd.Series):
        return info.to_dict()
    elif isinstance(info, dict):
        return info
    elif isinstance(info, pd.DataFrame):
        return info.iloc[0].to_dict()
    else:
        return {}

def get_combined_json(df_final, stock_code):
    # 获取元数据
    raw_metadata = get_stock_info_json(stock_code)
    metadata = rename_keys(raw_metadata)
    # 转换数据部分
    data_json = json.loads(df_final.to_json(orient='records', force_ascii=False))
    
    # 构建结果
    result = {
        "data": data_json,
        "metadata": metadata
    }
    
    # 返回 JSON 字符串
    return json.dumps(result, ensure_ascii=False, indent=2)

def rename_keys(data):
    """重命名键的函数"""
    # 定义需要重命名的键映射
    key_mapping = {
        "市盈率(动)": "动态市盈率",
        # 可以添加其他需要重命名的键
    }
    
    # 如果是字典，直接重命名
    if isinstance(data, dict):
        # 创建新字典
        new_data = {}
        for key, value in data.items():
            # 检查是否需要重命名
            new_key = key_mapping.get(key, key)
            # 递归处理嵌套值
            new_data[new_key] = rename_keys(value)
        return new_data
    
    # 如果是列表，递归处理每个元素
    if isinstance(data, list):
        return [rename_keys(item) for item in data]
    
    # 其他类型直接返回
    return data

# 修改主函数
if __name__ == '__main__':
    stock_code = sys.argv[1]
    start_date = sys.argv[2] if len(sys.argv) > 2 else None
    end_date = sys.argv[3] if len(sys.argv) > 3 else None
    # stock_code = '601077'
    # start_date='20250415'
    # end_date='20250607'
    # session = requests.Session()
    # session.verify = False
    # ef.utils.requests_session = session
    df_quote = ef.stock.get_quote_history(stock_code, beg=start_date, end=end_date, klt=101)

    # 获取资金流向数据（全量）
    df_bill = ef.stock.get_history_bill(stock_code)

    # 转换日期格式
    # df_quote['日期'] = pd.to_datetime(df_quote['日期'])
    # df_bill['日期'] = pd.to_datetime(df_bill['日期'])

    df_quote = calculate_indicators(df_quote)
    # 合并两个 DataFrame（inner join）
    # 关键修改：使用日期列合并，而不是索引
    df_final = pd.merge(
        df_quote,
        df_bill,
        on='日期',  # 使用日期列作为合并键
        how='inner',
        suffixes=('', '_bill')  # 为重复列添加后缀
    )
    
    # 删除重复列（带有_bill后缀的列）
    cols_to_drop = [col for col in df_final.columns if col.endswith('_bill')]
    df_final.drop(columns=cols_to_drop, inplace=True)

    # 计算指标
    
    # 将日期格式化为 "YYYY-MM-DD"
    # df_final['日期'] = df_final['日期'].dt.strftime('%Y-%m-%d')

    json_output = get_combined_json(df_final, stock_code)
    print(json_output)
    # 转换为 JSON 字符串并打印

    # print(df_final.to_json(orient='records', force_ascii=False))